/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2020-2021. All rights reserved.
 * Description: 集合通信all reduce算子融合
 * Author: wangcimeng
 * Create: 2020-5-12
 */

#include "hvd_all_reduce_fusion.h"
#include "hcom_ops_kernel_info_store.h"
#include "comm.h"
#include "hccl/hcom.h"
#include "hvd_ops_kernel_info_store.h"
#include "graph/debug/ge_attr_define.h"
#include "graph/utils/graph_utils.h"
#include "graph/utils/tensor_utils.h"
#include "graph/node.h"
using namespace std;

namespace hccl {
HvdAllReduceFusion::HvdAllReduceFusion()
{
}

HvdAllReduceFusion::~HvdAllReduceFusion()
{
}

HcclResult HvdAllReduceFusion::Run(ge::ComputeGraph& graph)
{
    HcclResult ret;
    std::vector<ge::NodePtr> fusionOps;
    ret = GetFusionOps(graph, fusionOps);
    CHK_PRT_RET(ret != HCCL_SUCCESS,
        HCCL_ERROR("[Run][HvdAllReduceFusion]graph[%s]: get fusion HcomAllReduce ops failed. ret[%d]",
            graph.GetName().c_str(), ret), ret);

    // The number of HcomAllReduce operator must be more than 1
    CHK_PRT_RET((fusionOps.size() <= 0), HCCL_INFO("NOT_CHANGED: the graph has %d no HcomAllReduce op.",
        fusionOps.size()), HCCL_SUCCESS);

    std::vector<u32> segmentIndex;
    ret = GetSplitStrategy(graph, fusionOps, segmentIndex);
    CHK_PRT_RET(ret != HCCL_SUCCESS,
        HCCL_ERROR("[Run][HvdAllReduceFusion]graph[%s]: get HcomAllReduce ops split strategy"
            "failed. ret[%d]", graph.GetName().c_str(), ret), ret);

    CHK_RET(RunFusionOps(graph, fusionOps, segmentIndex.size(), segmentIndex));

    HCCL_INFO("fuse HcomAllReduce op end");
    return HCCL_SUCCESS;
}

HcclResult HvdAllReduceFusion::GetFusionOps(ge::ComputeGraph& graph, std::vector<ge::NodePtr>& fusionOps)
{
    for (auto nodePtr : graph.GetDirectNode()) {
        if (!nodePtr) {
            HCCL_WARNING("HvdAllReduceFusion: null node exists.");
            continue;
        }
        auto opDescPtr = nodePtr->GetOpDesc();
        if (!opDescPtr) {
            HCCL_WARNING("HvdAllReduceFusion: desc of node[%s] is null.", nodePtr->GetName().c_str());
            continue;
        }
        if (opDescPtr->GetType() == HVD_KERNEL_OP_TYPE_ALLREDUCE) {
            fusionOps.push_back(nodePtr);
        }
    }
    HCCL_INFO("there are %d op[%s] to be fused in graph[%s].", fusionOps.size(),
        HVD_KERNEL_OP_TYPE_ALLREDUCE.c_str(), graph.GetName().c_str());
    return HCCL_SUCCESS;
}

HcclResult HvdAllReduceFusion::GetSplitStrategy(ge::ComputeGraph& graph, std::vector<ge::NodePtr>& fusionOps,
    std::vector<u32>& segmentIndex)
{
    std::vector<u32> segList;
    u32 gradientNum = fusionOps.size();
    u32 firstSegment = static_cast<u32>(gradientNum * HVD_GRADIENT_SEGMENT_INDEX_RATIO);
    segList.push_back(firstSegment);
    if (gradientNum >= (firstSegment + 2)) { // 2: 保证第二段大于0
        segList.push_back(gradientNum - 1);
    }
    segmentIndex = segList;
    std::string printStr;
    u32 baseIndex = 0;

    for (u32 i = 0; i < segmentIndex.size(); i++) {
        printStr.append("[");
        printStr.append(std::to_string(baseIndex));
        printStr.append(",");
        printStr.append(std::to_string(segmentIndex[i]));
        printStr.append("] ");
        baseIndex = segmentIndex[i] + 1;
    }
    HCCL_INFO("gradient segment result: segment num: %u, segment index list: %s ",   \
        segmentIndex.size(), printStr.c_str());

    return HCCL_SUCCESS;
}
}
